package SparkSQL

import org.apache.spark.sql.expressions.{Aggregator}
import org.apache.spark.sql.{Encoder, Encoders, SparkSession, functions}

/**
 * 强类型
 * UDAF是自定义函数
 * 他和UDF不一样的地方是UDF是一对一输出的
 * 但是UDAF是聚合函数（类似于sum、count）
 */
object UDAFTest02 {
  def main(args: Array[String]): Unit = {
    val spark = SparkSession.builder()
      .master("local[*]")
      .appName("UDAF")
      .getOrCreate()
    import spark.implicits._
    val df = spark.read
      .json("data/user.json")
    df.createOrReplaceTempView("user")
    val ds = df.as[User]

    //将UDAF变成查询的列对象
    val udafCol = new MyAvgAgeUDAF().toColumn
    ds.select(udafCol).show()
    spark.close()
  }


  /**
   * 使用var的原因：在Scala中样例类的数据的默认是无法更改的，
   * 使用var就是为了可以更改样例类的属性
   *
   * @param total
   * @param count
   */
  case class Buff(var total: Long, var count: Long)

  case class User(var username: String, var age: Long)

  /**
   * 继承:Aggregator(org.apache.spark.sql.expressions.Aggregator)
   * 定义泛型：
   * IN :输入的数据类型
   * BUF:缓冲区的数据类型
   * OUT:输出的数据类型
   * 重写方法:
   */
  class MyAvgAgeUDAF extends Aggregator[User, Buff, Long] {
    /**
     * 初始值或者是零值
     * 缓冲区的初始化
     *
     * @return
     */
    override def zero: Buff = {
      Buff(0l, 0l)
    }

    /**
     * 根据输入的数据来更新缓冲区的数据
     *
     * @param b
     * @param a
     * @return
     */
    override def reduce(b: Buff, a: User): Buff = {
      b.total += a.age
      b.count += 1
      b
    }

    /**
     * 合并缓冲区
     *
     * @param b1
     * @param b2
     * @return
     */
    override def merge(b1: Buff, b2: Buff): Buff = {
      b1.total += b2.total
      b1.count += b2.count
      b1
    }

    /**
     * 计算结果
     *
     * @param reduction
     * @return
     */
    override def finish(reduction: Buff): Long = (reduction.total / reduction.count)

    /**
     * 这是固定的写法，若是自定义的类那么就是：product
     * 缓冲区的编码操作
     *
     * @return
     */
    override def bufferEncoder: Encoder[Buff] = Encoders.product

    /**
     * 这也是固定的写法，若是scala存在的类（如long，int，string……）就是选择对应的即可
     * 输出的编码操作
     *
     * @return
     */
    override def outputEncoder: Encoder[Long] = Encoders.scalaLong
  }
}
